#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
TCP服务器协议实现
支持多客户端连接和样本数据推送
"""

import socket
import threading
import time
import logging
import json
import struct
from datetime import datetime
from typing import Dict, List, Optional, Callable
from PyQt6.QtCore import QThread, pyqtSignal

try:
    from src.utils.json_encoder import safe_json_dumps, safe_json_loads, encode_json_message, decode_json_message
except ImportError:
    try:
        from ..utils.json_encoder import safe_json_dumps, safe_json_loads, encode_json_message, decode_json_message
    except ImportError:
        # 如果导入失败，使用标准json模块
        def safe_json_dumps(data, **kwargs):
            return json.dumps(data, ensure_ascii=False, default=str, **kwargs)
        
        def safe_json_loads(json_str, **kwargs):
            return json.loads(json_str, **kwargs)
        
        def encode_json_message(data):
            json_str = safe_json_dumps(data)
            json_bytes = json_str.encode('utf-8')
            length_bytes = struct.pack('>I', len(json_bytes))
            return length_bytes + json_bytes
        
        def decode_json_message(data):
            if len(data) < 4:
                return None
            message_length = struct.unpack('>I', data[:4])[0]
            if len(data) < 4 + message_length:
                return None
            json_bytes = data[4:4 + message_length]
            json_str = json_bytes.decode('utf-8')
            return safe_json_loads(json_str)

logger = logging.getLogger(__name__)

class TCPClientConnection:
    """TCP客户端连接"""
    
    def __init__(self, client_socket: socket.socket, client_address: tuple, server):
        self.client_socket = client_socket
        self.client_address = client_address
        self.server = server
        self.connected = True
        self.last_heartbeat = time.time()
        self.client_id = f"{client_address[0]}:{client_address[1]}"
        
        # 设置socket选项
        self.client_socket.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
        self.client_socket.settimeout(1.0)
        
        logger.info(f"客户端连接建立: {self.client_id}")
    
    def send_data(self, data: bytes) -> bool:
        """发送数据到客户端"""
        try:
            if self.connected and self.client_socket:
                self.client_socket.send(data)
                logger.debug(f"向客户端 {self.client_id} 发送数据: {len(data)} 字节")
                return True
            return False
        except Exception as e:
            logger.error(f"向客户端 {self.client_id} 发送数据失败: {e}")
            self.connected = False
            return False
    
    def send_json(self, data: dict) -> bool:
        """发送JSON数据到客户端"""
        try:
            # 使用安全的JSON编码
            message_data = encode_json_message(data)
            return self.send_data(message_data)
        except Exception as e:
            logger.error(f"向客户端 {self.client_id} 发送JSON数据失败: {e}")
            return False
    
    def close(self):
        """关闭连接"""
        self.connected = False
        try:
            if self.client_socket:
                self.client_socket.close()
        except Exception as e:
            logger.error(f"关闭客户端连接失败: {e}")
        logger.info(f"客户端连接已关闭: {self.client_id}")

class TCPServerThread(QThread):
    """TCP服务器线程"""
    
    client_connected = pyqtSignal(str)  # 客户端连接信号
    client_disconnected = pyqtSignal(str)  # 客户端断开信号
    data_received = pyqtSignal(str, bytes)  # 数据接收信号 (client_id, data)
    
    def __init__(self, host: str = '0.0.0.0', port: int = 5000, max_clients: int = 10):
        super().__init__()
        self.host = host
        self.port = port
        self.max_clients = max_clients
        self.server_socket = None
        self.running = False
        self.clients: Dict[str, TCPClientConnection] = {}
        self.clients_lock = threading.Lock()
        
    def run(self):
        """运行服务器"""
        try:
            self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            self.server_socket.bind((self.host, self.port))
            self.server_socket.listen(self.max_clients)
            self.server_socket.settimeout(1.0)
            
            self.running = True
            logger.info(f"TCP服务器启动成功: {self.host}:{self.port}")
            
            while self.running:
                try:
                    client_socket, client_address = self.server_socket.accept()
                    
                    # 检查客户端数量限制
                    with self.clients_lock:
                        if len(self.clients) >= self.max_clients:
                            logger.warning(f"客户端数量已达上限 ({self.max_clients})，拒绝连接: {client_address}")
                            client_socket.close()
                            continue
                    
                    # 创建客户端连接
                    client_conn = TCPClientConnection(client_socket, client_address, self)
                    client_id = client_conn.client_id
                    
                    with self.clients_lock:
                        self.clients[client_id] = client_conn
                    
                    # 启动客户端处理线程
                    client_thread = threading.Thread(
                        target=self._handle_client,
                        args=(client_conn,),
                        daemon=True
                    )
                    client_thread.start()
                    
                    # 发送连接信号
                    self.client_connected.emit(client_id)
                    
                except socket.timeout:
                    # 超时，继续循环
                    continue
                except Exception as e:
                    if self.running:
                        logger.error(f"接受客户端连接失败: {e}")
                    break
                    
        except Exception as e:
            logger.error(f"TCP服务器启动失败: {e}")
        finally:
            self.stop()
    
    def _handle_client(self, client_conn: TCPClientConnection):
        """处理客户端连接"""
        client_id = client_conn.client_id
        
        try:
            while client_conn.connected and self.running:
                try:
                    # 接收数据
                    data = client_conn.client_socket.recv(4096)
                    if not data:
                        logger.info(f"客户端 {client_id} 断开连接")
                        break
                    
                    # 发送数据接收信号
                    self.data_received.emit(client_id, data)
                    
                except socket.timeout:
                    # 超时，继续循环
                    continue
                except Exception as e:
                    logger.error(f"处理客户端 {client_id} 数据失败: {e}")
                    break
                    
        except Exception as e:
            logger.error(f"客户端 {client_id} 处理异常: {e}")
        finally:
            # 清理客户端连接
            with self.clients_lock:
                if client_id in self.clients:
                    del self.clients[client_id]
            
            client_conn.close()
            self.client_disconnected.emit(client_id)
    
    def stop(self):
        """停止服务器"""
        self.running = False
        
        # 关闭所有客户端连接
        with self.clients_lock:
            for client_conn in self.clients.values():
                client_conn.close()
            self.clients.clear()
        
        # 关闭服务器socket
        if self.server_socket:
            try:
                self.server_socket.close()
            except Exception as e:
                logger.error(f"关闭服务器socket失败: {e}")
        
        logger.info("TCP服务器已停止")
    
    def get_client_count(self) -> int:
        """获取客户端数量"""
        with self.clients_lock:
            return len(self.clients)
    
    def get_client_list(self) -> List[str]:
        """获取客户端列表"""
        with self.clients_lock:
            return list(self.clients.keys())
    
    def send_to_client(self, client_id: str, data: bytes) -> bool:
        """向指定客户端发送数据"""
        with self.clients_lock:
            # 如果client_id为None，广播给所有客户端
            if client_id is None:
                if not self.clients:
                    return False
                success_count = 0
                for client_conn in self.clients.values():
                    if client_conn.send_data(data):
                        success_count += 1
                return success_count > 0
            # 如果指定了client_id，发送给特定客户端
            elif client_id in self.clients:
                return self.clients[client_id].send_data(data)
        return False
    
    def send_json_to_client(self, client_id: str, data: dict) -> bool:
        """向指定客户端发送JSON数据"""
        with self.clients_lock:
            # 如果client_id为None，广播给所有客户端
            if client_id is None:
                if not self.clients:
                    return False
                success_count = 0
                for client_conn in self.clients.values():
                    if client_conn.send_json(data):
                        success_count += 1
                return success_count > 0
            # 如果指定了client_id，发送给特定客户端
            elif client_id in self.clients:
                return self.clients[client_id].send_json(data)
        return False
    
    def broadcast_data(self, data: bytes) -> int:
        """广播数据到所有客户端"""
        success_count = 0
        with self.clients_lock:
            for client_conn in self.clients.values():
                if client_conn.send_data(data):
                    success_count += 1
        return success_count
    
    def broadcast_json(self, data: dict) -> int:
        """广播JSON数据到所有客户端"""
        success_count = 0
        with self.clients_lock:
            for client_conn in self.clients.values():
                if client_conn.send_json(data):
                    success_count += 1
        return success_count

class TCPServer:
    """TCP服务器管理器"""
    
    def __init__(self, host: str = '0.0.0.0', port: int = 5000, max_clients: int = 10):
        self.host = host
        self.port = port
        self.max_clients = max_clients
        self.server_thread = None
        self.is_running = False
        
    def start(self) -> bool:
        """启动服务器"""
        try:
            if self.is_running:
                logger.warning("TCP服务器已在运行")
                return True
            
            self.server_thread = TCPServerThread(self.host, self.port, self.max_clients)
            self.server_thread.start()
            self.is_running = True
            
            logger.info(f"TCP服务器启动成功: {self.host}:{self.port}")
            return True
            
        except Exception as e:
            logger.error(f"启动TCP服务器失败: {e}")
            return False
    
    def stop(self):
        """停止服务器"""
        if self.server_thread:
            self.server_thread.stop()
            self.server_thread.wait()
            self.server_thread = None
        
        self.is_running = False
        logger.info("TCP服务器已停止")
    
    def is_server_running(self) -> bool:
        """检查服务器是否运行"""
        return self.is_running and self.server_thread and self.server_thread.isRunning()
    
    def get_client_count(self) -> int:
        """获取客户端数量"""
        if self.server_thread:
            return self.server_thread.get_client_count()
        return 0
    
    def get_client_list(self) -> List[str]:
        """获取客户端列表"""
        if self.server_thread:
            return self.server_thread.get_client_list()
        return []
    
    def send_to_client(self, client_id: str, data: bytes) -> bool:
        """向指定客户端发送数据"""
        if self.server_thread:
            return self.server_thread.send_to_client(client_id, data)
        return False
    
    def send_json_to_client(self, client_id: str, data: dict) -> bool:
        """向指定客户端发送JSON数据"""
        if self.server_thread:
            return self.server_thread.send_json_to_client(client_id, data)
        return False
    
    def broadcast_data(self, data: bytes) -> int:
        """广播数据到所有客户端"""
        if self.server_thread:
            return self.server_thread.broadcast_data(data)
        return 0
    
    def broadcast_json(self, data: dict) -> int:
        """广播JSON数据到所有客户端"""
        if self.server_thread:
            return self.server_thread.broadcast_json(data)
        return 0
    
    def set_callbacks(self, on_client_connected: Callable = None, 
                     on_client_disconnected: Callable = None,
                     on_data_received: Callable = None):
        """设置回调函数"""
        if self.server_thread:
            if on_client_connected:
                self.server_thread.client_connected.connect(on_client_connected)
            if on_client_disconnected:
                self.server_thread.client_disconnected.connect(on_client_disconnected)
            if on_data_received:
                self.server_thread.data_received.connect(on_data_received) 